import logging
import os
import base64
from typing import List, Dict, Optional
from pathlib import Path

# 创建日志记录器
logger = logging.getLogger(__name__)


def get_datasets_list(data_dir: str = "data") -> List[Dict]:
    """
    获取数据集列表，支持2D图像数据集和3D医学数据集
    
    Args:
        data_dir: 数据目录路径
        
    Returns:
        包含数据集信息的列表
    """
    logger.info(f"开始获取数据集列表，数据目录: {data_dir}")
    
    datasets = []
    
    # 检查数据目录是否存在
    if not os.path.exists(data_dir):
        logger.warning(f"数据目录不存在: {data_dir}")
        return datasets
    
    # 检查是否为目录
    if not os.path.isdir(data_dir):
        logger.warning(f"数据路径不是目录: {data_dir}")
        return datasets
    
    logger.debug(f"数据目录存在，开始遍历: {data_dir}")
    
    # 遍历数据目录下的所有子目录
    for dataset_name in os.listdir(data_dir):
        dataset_path = os.path.join(data_dir, dataset_name)
        
        # 只处理目录
        if os.path.isdir(dataset_path):
            logger.debug(f"发现数据集目录: {dataset_name}")
            
            # 计算图像文件数量
            image_count = 0
            medical_image_count = 0
            for file in os.listdir(dataset_path):
                # 2D图像文件
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
                    image_count += 1
                # 3D医学图像文件
                elif file.lower().endswith(('.nii', '.nii.gz', '.dicom')):
                    medical_image_count += 1
            
            # 确定数据集类型
            dataset_type = "2d"  # 默认为2D
            if medical_image_count > 0:
                dataset_type = "3d"  # 如果有3D医学图像文件，则为3D
            elif image_count > 0:
                dataset_type = "2d"  # 如果只有2D图像文件，则为2D
            
            dataset_info = {
                # "id": dataset_name,
                "name": dataset_name,
                "type": dataset_type,
                "image_count": image_count,
                "medical_image_count": medical_image_count,
                "path": dataset_path
            }
            
            datasets.append(dataset_info)
            logger.debug(f"添加数据集: {dataset_info}")
    
    logger.info(f"数据集列表获取完成，共找到 {len(datasets)} 个数据集")
    return datasets


def get_dataset_info(dataset_id: str, data_dir: str = "data") -> Optional[Dict]:
    """
    获取特定数据集的详细信息，支持2D图像数据集和3D医学数据集
    
    Args:
        dataset_id: 数据集ID
        data_dir: 数据目录路径
        
    Returns:
        数据集详细信息，如果未找到则返回None
    """
    logger.info(f"开始获取数据集详情，数据集ID: {dataset_id}, 数据目录: {data_dir}")
    
    dataset_path = os.path.join(data_dir, dataset_id)
    
    # 检查数据集目录是否存在
    if not os.path.exists(dataset_path):
        logger.warning(f"数据集目录不存在: {dataset_path}")
        return None
    
    if not os.path.isdir(dataset_path):
        logger.warning(f"数据集路径不是目录: {dataset_path}")
        return None
    
    logger.debug(f"数据集目录存在: {dataset_path}")
    
    # 获取图像文件列表
    images = []
    medical_images = []
    try:
        for file in os.listdir(dataset_path):
            # 处理2D图像文件
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
                # 使用文件名（不含扩展名）作为图像ID
                image_id = Path(file).stem
                image_info = {
                    "id": image_id,
                    "filename": file,
                    "path": os.path.join(dataset_path, file),
                    "type": "2d"
                }
                images.append(image_info)
                logger.debug(f"添加2D图像信息: {image_info}")
            # 处理3D医学图像文件
            elif file.lower().endswith(('.nii', '.nii.gz', '.dicom')):
                # 使用文件名（不含扩展名）作为图像ID
                image_id = Path(file).stem
                # 特殊处理.nii.gz扩展名
                if file.lower().endswith('.nii.gz'):
                    image_id = Path(file).stem.split('.')[0]  # 移除.nii部分，只保留基础文件名
                image_info = {
                    "id": image_id,
                    "filename": file,
                    "path": os.path.join(dataset_path, file),
                    "type": "3d"
                }
                medical_images.append(image_info)
                logger.debug(f"添加3D医学图像信息: {image_info}")
    except Exception as e:
        logger.error(f"遍历数据集目录时发生错误: {dataset_path}, 错误: {str(e)}")
        return None
    
    # 合并图像列表
    all_images = images + medical_images
    
    # 确定数据集类型
    dataset_type = "2d"  # 默认为2D
    if len(medical_images) > 0:
        dataset_type = "3d"  # 如果有3D医学图像文件，则为3D
    elif len(images) > 0:
        dataset_type = "2d"  # 如果只有2D图像文件，则为2D
    
    dataset_info = {
        "id": dataset_id,
        "name": dataset_id,
        "type": dataset_type,
        "image_count": len(images),
        "medical_image_count": len(medical_images),
        "path": dataset_path,
        "images": all_images
    }
    
    logger.info(f"数据集详情获取完成，数据集ID: {dataset_id}, 2D图像数量: {len(images)}, 3D医学图像数量: {len(medical_images)}")
    return dataset_info


def read_image_base64(image_path: str) -> Optional[str]:
    """
    读取图像文件并返回base64编码
    
    Args:
        image_path: 图像文件路径
        
    Returns:
        图像的base64编码字符串，如果读取失败则返回None
    """
    logger.debug(f"开始读取图像文件并转换为base64，文件路径: {image_path}")
    
    try:
        # 检查文件是否存在
        if not os.path.exists(image_path):
            logger.warning(f"图像文件不存在: {image_path}")
            return None
            
        # 读取图像文件
        with open(image_path, "rb") as image_file:
            # 将图像数据编码为base64
            encoded_string = base64.b64encode(image_file.read()).decode('utf-8')
            logger.debug(f"图像文件读取并转换完成，文件路径: {image_path}")
            return encoded_string
    except Exception as e:
        logger.error(f"读取图像文件时出错，文件路径: {image_path}, 错误: {str(e)}")
        return None


def get_image_info(image_id: str, dataset_id: str, data_dir: str = "data") -> Optional[Dict]:
    """
    获取特定图像的信息，支持2D图像和3D医学图像
    
    Args:
        image_id: 图像ID
        dataset_id: 数据集ID
        data_dir: 数据目录路径
        
    Returns:
        图像信息，如果未找到则返回None
    """
    logger.info(f"开始获取图像信息，图像ID: {image_id}, 数据集ID: {dataset_id}, 数据目录: {data_dir}")
    
    dataset_path = os.path.join(data_dir, dataset_id)
    
    # 检查数据集目录是否存在
    if not os.path.exists(dataset_path):
        logger.warning(f"数据集目录不存在: {dataset_path}")
        return None
    
    if not os.path.isdir(dataset_path):
        logger.warning(f"数据集路径不是目录: {dataset_path}")
        return None
    
    logger.debug(f"数据集目录存在，开始查找图像: {dataset_path}")
    
    # 查找匹配的图像文件
    try:
        for file in os.listdir(dataset_path):
            logger.info(f"检查文件: {file}")
            # 检查文件是否匹配image_id
            file_matches = False
            file_type = "unknown"
            
            # 检查2D图像文件
            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
                if Path(file).stem == image_id:
                    file_matches = True
                    file_type = "2d"
            # 检查3D医学图像文件
            elif file.lower().endswith(('.nii', '.nii.gz', '.dicom')):
                # 特殊处理.nii.gz扩展名
                if file.lower().endswith('.nii.gz'):
                    if Path(file).stem.split('.')[0] == image_id:
                        file_matches = True
                        file_type = "3d"
                else:
                    if Path(file).stem == image_id:
                        file_matches = True
                        file_type = "3d"
            
            # 如果文件匹配，返回图像信息
            if file_matches:
                image_info = {
                    "id": image_id,
                    "filename": file,
                    "path": os.path.join(dataset_path, file),
                    "type": file_type
                }
                logger.info(f"找到匹配的图像信息: {image_info}")
                return image_info
    except Exception as e:
        logger.error(f"遍历数据集目录查找图像时发生错误: {dataset_path}, 错误: {str(e)}")
        return None
    
    logger.warning(f"未找到匹配的图像，图像ID: {image_id}, 数据集ID: {dataset_id}")
    return None